/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transl;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.TreeMap;
import java.util.TreeSet;

import alloy.ast.ASTDepthFirstVisitor;
import alloy.ast.DeclIter;
import alloy.ast.FindCommand;
import alloy.ast.Formula;
import alloy.ast.HasType;
import alloy.ast.Leaf;
import alloy.ast.LeafExpr;
import alloy.ast.LeafId;
import alloy.ast.Module;
import alloy.ast.MultiplicityExpr;
import alloy.ast.Node;
import alloy.ast.SigExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.type.BasicType;
import alloy.type.RelationType;
import alloy.type.SigType;
import alloy.util.Dbg;

/**
 * Gathers various information about the translatable AST.
 * This information is then stored in one place (this visitor object)
 * for quick access.
 */
public class ASTInfoVisitor extends ASTDepthFirstVisitor implements ASTInfo {
    /** The Alloy command we were analyzing */
    private FindCommand _command;

    /**
     * The root {@link Formula} of the translatable AST.  Stored here in order to have
     * all the information about the formula in one place.
     */
    private Formula _formula;

    /** {@link Map} from {@link Node} to its parent. */
    private Map /* Node->Node */ _node2parent = new HashMap();

    /**
     * {@link Map} from {@link LeafId} to the {@link RelationType} of the
     * leaves with that {@link LeafId}.  Only relation-valued leaves
     * (relations and quantified variables) have entries in this map.
     */
    private Map /* LeafId->RelationType */ _leafType = new HashMap();

    /**
     * {@link Map} from {@link LeafId} of a quantified {@link Variable} or a relation
     * to its declared {@link MultiplicityExpr multiplicity}.
     */
    private Map /* LeafId->MultiplicityExpr */ _leafMult = new HashMap();

    /**
     * {@link Map} from {@link LeafId} to a representative leaf with that {@link LeafId}.
     * Note that for a quantified variable, the representative leaf will be
     * a {@link VariableExpr} (with {@link VariableExpr#isQuantified} set to
     * <code>true</code>), rather than a {@link Variable} (which is not a {@link Leaf}).
     */
    private Map /* LeafId->Leaf */ _leafRep = new HashMap();

    /**
     * For each quantified variable, the {@link VarCreator} node that defined it.
     * {@link Map} from {@link LeafId} of the quantified variable to the {@link VarCreator}
     * node.
     */
    private Map /* LeafId->VarCreator */ _var2creator = new HashMap();

    /**
     * For each quantified variable, the {@link Variable} object.
     * {@link Map} from the {@link LeafId} of the quantified variable to the {@link Variable}
     * object of the quantified variable.
     */
    private Map /* LeafId->Variable */ _leaf2var = new HashMap();

    /**
     * {@link Map} from each {@link VarCreator} to the {@link Set} of {@link LeafId}s of
     * quantified variables declared in that {@link VarCreator}.
     */
    private Map /* VarCreator->List[LeafId] */ _creator2vars = new HashMap();

    /** {@link Set} of {@link LeafId}s of relations used in the formula */
    private Set /* of LeafId */ _relationLeafIds = new TreeSet();

    /** {@link Set} of all {@link BasicType}s used in the formula */
    private Set /* of BasicType */ _allBasicTypes = new TreeSet();

    /**
     * The singleton {@link BasicType} object representing the special Int basic type.
     * null if Int is not used in the model.
     */
    private BasicType _intType;

    /**
     * {@link Map} from {@link LeafId} of each relation to a {@link LeafExpr}
     * for that relation: one {@link Map.Entry} per relation.
     */
    private Map /* LeafId->LeafExpr */ _relationReps;

    /**
     * Map from Node to Integer ID of the node.  Each node of the
     * translatable AST is deterministically assigned an integer ID
     * that depends only on the input model (i.e. will be the same on
     * different runs, different machines etc).
     */
    private Map /* Node->Integer */ _node2id = new HashMap();
    
    /** The next ID to give to a node of the translatable AST */
    private int _nextNodeID = 1;
    
    /////////////
    // Methods //
    /////////////
	
	public ASTInfoVisitor(FindCommand command_) {
	    _command = command_;
	    this._formula = command_.formula;
	    _leafMult.putAll(command_.leafIdToMult);
	}
    
    public void visit(Node node_) {
	// give the node a deterministic integer ID
	if (!_node2id.containsKey(node_)) {
	    int thisNodeID = _nextNodeID++;
	    _node2id.put(node_, new Integer(thisNodeID));
	    
	}
	//System.out.println("ID " + thisNodeID + ": " + node_.nodeString());
	
	if (node_ instanceof HasType)
	    _recordBasicTypes(((HasType)node_).getType().toRelationType());
	
	if (node_ instanceof Leaf) {
	    LeafId leafId = ((Leaf)node_).getLeafId();
	    Dbg.chk(leafId);
	    if (node_ instanceof VariableExpr &&
		((VariableExpr)node_).isQuantified) {
		Dbg.chk(_var2creator.containsKey(leafId), node_.fullString());
	    }
	    if (node_ instanceof HasType)
		_recordLeafType(leafId, ((HasType)node_).getType().toRelationType());
	    _leafRep.put(leafId, node_);
	    //System.out.println("MAPPING " + leafId + " TO " + node_.nodeString());
	    if (node_ instanceof SigExpr ||
		(node_ instanceof VariableExpr &&
		 !((VariableExpr)node_).isQuantified))
		_relationLeafIds.add(leafId);
	} else {  // not a leaf
	    if (node_ instanceof VarCreator) {
        List /* of LeafId */ declaredVars = new ArrayList();
        for (DeclIter declIter = new DeclIter((VarCreator)node_); declIter.hasNext();) {
            declIter.next();
            LeafId leafId = declIter.getLeafId();
            Dbg.chk(leafId);
            Variable var = declIter.getVariable();
            MultiplicityExpr mult = declIter.getMultiplicity();
            declaredVars.add(leafId);
            Dbg.chk(!_var2creator.containsKey(leafId));
            Dbg.chk(!_leaf2var.containsKey(leafId));
            if (!_leafMult.containsKey(leafId)) {
            //System.out.println("no multiplicity for var: " + var +
            //         " in formula " + node_);
            _leafMult.put(leafId, mult);
            }
            Dbg.chk(_leafMult.containsKey(leafId));
            _var2creator.put(leafId, node_);
            _leaf2var.put(leafId, var);
            _recordLeafType(leafId, declIter.getType().toRelationType());
            //_leafMult.put(leafId, mult);
        }
        _creator2vars.put(node_, Collections.unmodifiableList(declaredVars));
        }

        // visit the children
        for (int i=0; i<node_.numChildren(); i++) {
        Node child = node_.childAt(i);
        if (!(node_ instanceof VarCreator &&
              child == ((VarCreator)node_).getDecls())) {
            _node2parent.put(child, node_);
            child.applyVisitor(this);
        }
        }   // loop over children
    }   // not a leaf
    }  // visit(Node)

    /**
     * Record the type of leaves with the given {@link LeafId}.  If we already have a type
     * for this LeafId, make sure the new type is the same.
     */
    private void _recordLeafType(LeafId leafId_, RelationType relType_) {
    RelationType oldType = (RelationType)_leafType.get(leafId_);
    if (oldType != null)
        Dbg.chk(oldType.equals(relType_));
    else {
        _leafType.put(leafId_, relType_);
        _recordBasicTypes(relType_);
    }
    }

    /**
     * Record all {@link BasicType}s in the given {@link RelationType}.
     */
    private void _recordBasicTypes(RelationType relType_) {
    for (int i=0; i<relType_.numBasicTypes(); i++)
        _recordBasicType(relType_.getBasicTypeAt(i));
    }

    /**
     * Record the basic type in the list of all basic types;
     * if this basic type is the special Int type, record a reference to it.
     */
    private void _recordBasicType(BasicType basicType_) {
	//System.out.println("recording basic type: " + basicType_);
	_allBasicTypes.add(basicType_);
	if (((SigType) basicType_).getSignature() == 
	    ((Module)_command.getParent().getParent()).getModuleInfo().getInt()) {
	    _intType = basicType_;
	}
    }

    /////////////////////////////////////////////////////////////////////////////

    /** Return the root of the translatable AST */
    public Formula getFormula() { return _formula; }

    /** Return the command we are analyzing */
    public FindCommand getCommand() { return _command; }

    /**
     * Return the parent of the given {@link Node} in the translatable AST,
     * or <code>null</code> if 'node_' is the root of the translatable AST.
     */
    public Node getParent(Node node_) { return (Node)_node2parent.get(node_); }

    /**
     * @return <code>true</code> iff anc_ is an ancestor of node_.
     *         Note that a node is <em>not</em> its own ancestor.
     */
    public boolean isAncestor(Node anc_, Node node_) {
    Node parent = getParent(node_);
    return parent!=null &&
        (parent == anc_ ||
         isAncestor(anc_, parent));
    }

    /** Test whether the given {@link LeafId} is for a quantified variable. */
    public boolean isQuantVar(LeafId leafId_) { return _var2creator.containsKey(leafId_); }

    /**
     * @return <code>true</code> iff anc_ is an ancestor of node_.
     *         Note that a node is <em>not</em> its own ancestor.
     */
    public void printAncestors(Node node_) {
    Node parent = getParent(node_);
    if (parent != null) {
        //System.out.println(parent.fullString());
        printAncestors(parent);
    }
    }

    /**
     * @return the {@link VarCreator} that defined the quantified variable with this
     * {@link LeafId}.
     */
    public VarCreator getVarCreator(LeafId leafId_) {
    return (VarCreator)_var2creator.get(leafId_);
    }

    /**
     * @return all {@link BasicType}s used in the formula.
     */
    public Set /* of BasicType */ getBasicTypes() { return Collections.unmodifiableSet(_allBasicTypes); }

    /**
     * Return the {@link BasicType} object representing the special Int basic type,
     * or null if Int is not used in the model.
     */
    public BasicType getIntType() { return _intType; }

    /**
     * @return {@link Set} of {@link LeafId}s of relations.
     * @see RelationAllocator
     */
    public Set /* of LeafId */ getRelations() { return _relationLeafIds; }

    /** @return the {@link RelationType} of leaves with the given {@link LeafId}. */
    public RelationType getLeafType(LeafId leafId_) {
    return (RelationType)Dbg.chk(_leafType.get(leafId_));
    }

    /** @return the {@link Variable} object with the given {@link LeafId} */
    public Variable getVariable(LeafId leafId_) { return (Variable)Dbg.chk(_leaf2var.get(leafId_)); }


    /**
     * @return the {@link MultiplicityExpr multiplicity} of quantified variable or relation with the
     * given {@link LeafId}
     */
    public MultiplicityExpr getLeafMultiplicity(LeafId leafId_) {
    return (MultiplicityExpr)Dbg.chk(_leafMult.get(leafId_));
    }

    /**
     * Return a representative {@link Leaf} with the given {@link LeafId}.
     */
    public Leaf getLeafRep(LeafId leafId_) { 
	
	return (Leaf)Dbg.chk(_leafRep.get(leafId_)); 

    }

    /**
     * Return the {@link Set} of {@link LeafId}s of quantified variables
     * declared at the given given {@link Node} (empty unless the node is a
     * {@link VarCreator}).
     */
    public List /* of LeafId */ getVarsDeclaredAt(Node node_) {
    List declaredIds = (List)_creator2vars.get(node_);
    return declaredIds==null ? Collections.EMPTY_LIST : declaredIds;
    }

    /**
     * @return {@link Map} from {@link LeafId} of each relation to
     * a {@link LeafExpr} for that relation.
     */
    public Map /* LeafId->LeafExpr */ getRelationReps() {
    if (_relationReps == null) {
        _relationReps = new TreeMap();
        for (Iterator relIdIter = getRelations().iterator(); relIdIter.hasNext();) {
        LeafId relId = (LeafId)relIdIter.next();
        _relationReps.put(relId, getLeafRep(relId));
        }
    }
    return _relationReps;
    }

    /**
     * Return an integer identifying the given node of the translatable
     * AST.  Each node of the translatable AST is deterministically assigned
     * an integer ID that depends only on the input model (i.e. will be the
     * same on different runs, different machines etc).
     */
    public int getNodeID(Node node_) {
    return ((Integer)_node2id.get(node_)).intValue();
    }

    /**
     * Test whether a given leaf is a skolem constant.
     */
    public boolean isSkolemConstant(LeafId leafId_) {
    return _command.skolemIds.contains(leafId_);
    }
}  // class ASTInfoVisitor
